package com.example.endusation.rag;


import com.example.endusation.neo4j.KnowledgeGraphService;
import com.example.endusation.ollama.OllamaDirectClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.util.*;

@Service
public class DiscreteMathRagService {
    @Autowired
    private KnowledgeGraphService knowledgeGraphService;

    @Autowired
    private OllamaDirectClient ollamaClient;

    @Value("${ollama.model}")
    private String modelName;
    private final String[] Entity = {"离散数学", "集合", "关系", "函数", "图", "树", "命题", "谓词", "二元关系", "笛卡儿积", "闭包", "单射", "满射", "双射", "逆函数", "无向图", "有向图", "欧拉图", "无向树", "有向树"};


    /**
     * 离散数学智能问答主流程
     */
    public String answerQuestion(String question,List<String> tokens) {

        // 1. 从知识图谱获取相关概念和定理
        List<Collection<Map<String, Object>>> list = new ArrayList<>();
        for (String n : tokens) {
            Collection<Map<String, Object>> kgInfo = knowledgeGraphService.queryConceptLayers(n, 1);
            list.add(kgInfo);
        }
        String neo4jData = "";
        for (Collection<Map<String, Object>> collection : list) {
            if (collection.size() == 0) {
                continue;
            }
            Iterator<Map<String, Object>> iterator = collection.iterator();
            // 遍历元素
            while (iterator.hasNext()) {
                Map<String, Object> element = iterator.next(); // 获取下一个元素
                neo4jData += (String) element.get("source") + element.get("relation") + element.get("target") + ";";
            }
        }

        // 2. 从Milvus获取相关文档片段


        //3. 融合检索结果，构建提示词
        String prompt = buildPrompt(question, neo4jData);

        // 4. 调用同义千问模型生成回答
        return generateAnswer(prompt);

    }

    /**
     * 构建离散数学领域的专业提示词
     */
    private String buildPrompt(String question, String neo4jData) {
        StringBuilder promptBuilder = new StringBuilder();

        promptBuilder.append(neo4jData+"\n");
        // 系统提示
        promptBuilder.append("请结合我给你的上文回答下面这个问题\n");
        promptBuilder.append("问题:"+question);
        return promptBuilder.toString();
    }

    /**
     * 调用Ollama中的同义千问模型生成回答
     */
    private String generateAnswer(String prompt) {
        return ollamaClient.chat(modelName, prompt);
    }

    private String[] getKnoledge(String question) {
        List<String> matches = new ArrayList<>();

        // 遍历所有实体，检查是否存在于问题中
        for (String entity : Entity) {
            // 使用contains方法检查是否包含
            if (question != null && question.contains(entity)) {
                matches.add(entity);
            }
        }

        // 将List转换为数组返回
        return matches.toArray(new String[0]);
    }
    public String answerQuestionJudge(String question) {
        if(question.contains("视频")){
            return "1";
        } else if(question.contains("图谱")){
            return "2";
        } else if (question.contains("资料")){
            return "3";
        } else {
            return "4";
        }
    }
}

